import os
import copy
import json
from pathlib import Path

import torch


class Trainer:
    def __init__(self, config, agents):
        self.config = config
        self.agents = agents
        self.results = dict()
        torch.set_num_threads(1)

    def train_agents(self):
        for _, agent_class in enumerate(self.agents):
            self.train_agent(agent_class)
        self.save_data_results('../results/data/')

    def train_agent(self, agent_class):
        agent_results = []
        agent_name = agent_class.agent_name
        seeds = self.config['seed']
        for agent_round in range(len(seeds)):
            agent_config = copy.deepcopy(self.config)
            agent_config['seed'] = seeds[agent_round]
            print(f"AGENT:{agent_name} ROUND:{agent_round} SEED:{agent_config['seed']}")
            agent = agent_class(agent_config)
            eval, avg_eval = agent.run()
            agent_results.append([eval, avg_eval, agent.get_score_required_to_win(), agent_config['seed']])
            print('-' * 80)
            print('-' * 80)
        self.results[agent_name] = agent_results

    def save_data_results(self, directory):
        env_name = self.config['env_name']
        for t in self.results.items():
            agent_result = dict()
            agent_result[t[0]] = t[1]
            path = Path(f'{directory}{env_name}_{t[0]}.json')
            json_str = json.dumps(agent_result, indent=4)
            path.write_text(json_str)

